#include <Vector.hpp>
#include <algorithm>
#include <iostream>
#include <numeric>
#include <queue>
#include <vector>

using namespace std;

class Solution
{
  public:
    vector<long long> findMaxSum(vector<int> &nums1, vector<int> &nums2, int k)
    {
        vector<tuple<int, int, int>> nums;
        for (int i = 0; i < nums1.size(); i++)
        {
            nums.emplace_back(nums1[i], i, nums2[i]);
        }
        sort(nums.begin(), nums.end(),
             [](const tuple<int, int, int> &a, const tuple<int, int, int> &b) { return get<0>(a) < get<0>(b); });

        priority_queue<int, vector<int>, greater<int>> ps;
        vector<long long> result(nums1.size());
        int maxsum = 0;
        for (int i = 0; i < nums.size();)
        {
            int start = i;
            int val = get<0>(nums[i]);
            while (i < nums.size() && get<0>(nums[i]) == val)
            {
                result[get<1>(nums[i])] = maxsum;
                i++;
            }

            for (int j = start; j < i; j++)
            {
                int num2 = get<2>(nums[j]);
                if (num2 > 0)
                {
                    maxsum += num2;
                    ps.push(num2);
                }
                if (ps.size() > k)
                {
                    maxsum -= ps.top();
                    ps.pop();
                }
            }
        }
        return result;
    }
};
int main()
{
    cout << "Code3478" << endl;
    Solution s;
    vector<int> nums1 = {4, 2, 1, 5, 3};
    vector<int> nums2 = {10, 20, 30, 40, 50};
    PrintVector(s.findMaxSum(nums1, nums2, 2));
    return 0;
}